# Copyright 2021-2023 @ Shenzhen Bay Laboratory &
#                       Peking University &
#                       Huawei Technologies Co., Ltd
#
# This code is a part of Cybertron package.
#
# The Cybertron is open-source software based on the AI-framework:
# MindSpore (https://www.mindspore.cn/)
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""
Aggregator for interaction layer
For Edge vector
"""

from typing import Union

import mindspore as ms
from mindspore import Tensor
from mindspore import nn
from mindspore import ops
from mindspore.ops import functional as F
from mindspore.common.initializer import initializer

from mindsponge.function import get_integer, get_arguments

from ..layer import Dense
from ..base import SoftmaxWithMask
from ..base import MultiheadAttention

__all__ = []

_AGGREGATOR_BY_KEY = dict()

def _aggregator_register(*aliases):
    """Return the alias register."""
    def alias_reg(cls):
        name = cls.__name__
        name = name.lower()
        if name not in _AGGREGATOR_BY_KEY:
            _AGGREGATOR_BY_KEY[name] = cls

        for alias in aliases:
            if alias not in _AGGREGATOR_BY_KEY:
                _AGGREGATOR_BY_KEY[alias] = cls

        return cls

    return alias_reg

class EdgeAggregator(nn.Cell):

    def __init__(self, dim: int, axis: tuple = (-3, -2), **kwargs):
        super().__init__()
        self._kwargs = kwargs

        self.dim = get_integer(dim)
        self.axis = axis
    
    def construct(self, 
                  edge_vec: Tensor, 
                  edge_mask: Tensor = None,
                  num_edges: Tensor = None,
                  **kwargs):
        
        return NotImplementedError

    def __str__(self):
        return 'EdgeAggregator<>'

@_aggregator_register('edge_sum')
class EdgeTensorSummation(EdgeAggregator):

    def __init__(self,
                 dim: int = None,
                 axis: tuple = (-3, -2),
                 **kwargs):

        super().__init__(dim=dim,
                       axis=axis,)
        
        self._kwargs = get_arguments(locals(), kwargs)
    
    def construct(self,
                  edge_vec: Tensor,
                  edge_mask: Tensor = None,
                  num_edges: Tensor = None,
                  edge_cutoff: Tensor = None,
                  ):
         
        if edge_cutoff is not None:
            edge_cutoff = F.expand_dims(edge_cutoff, axis=-1)
            if edge_mask is not None:
                edge_mask = F.expand_dims(edge_mask, axis=-1)
                # (B, A, A, 1) * (B, A, A, 1)
                edge_cutoff = F.mul(edge_cutoff, F.cast(edge_mask, ms.float32))
                # (B, A, A, F) * (B, A, A, 1)
                edge_vec = F.mul(edge_vec, edge_cutoff)
            else:
                # (B, A, A, F) * (B, A, A, 1)
                # print("edge_cutoff:\n", edge_cutoff)
                edge_vec = F.mul(edge_vec, edge_cutoff)
        elif edge_mask is not None:
            # (B, A, A, F) * (B, A, A, 1)
            edge_mask = F.expand_dims(edge_mask, axis=-1)
            edge_vec = F.mul(edge_vec, F.expand_dims(edge_mask, axis=-1))
        else:
            edge_vec = edge_vec
        # (B, A, A, F) -> (B, F)
        agg = F.reduce_sum(edge_vec, self.axis)
        
        return agg
    
    def __str__(self):
        return "EdgeTensorSummation<>"
    
@_aggregator_register('edge_mean')
class EdgeTensorMean(EdgeAggregator):
    
    def __init__(self,
                dim: int = None,
                axis: tuple = (-3, -2),
                **kwargs):

        super().__init__(dim=dim,
                         axis=axis,)
        
        self._kwargs = get_arguments(locals(), kwargs)
    
    def construct(self,
                edge_vec: Tensor,
                edge_mask: Tensor = None,
                num_edges: Tensor = None,
                edge_cutoff: Tensor = None,
                ):
        
        if edge_cutoff is not None:
            edge_cutoff = F.expand_dims(edge_cutoff, axis=-1)
            if edge_mask is not None:
                edge_mask = F.expand_dims(edge_mask, axis=-1)
                # (B, A, A, 1) * (B, A, A, 1)
                edge_cutoff = F.mul(edge_cutoff, F.cast(edge_mask, ms.float32))
                # (B, A, A, F) * (B, A, A, 1)
                edge_vec = F.mul(edge_vec, edge_cutoff)
            else:
                # (B, A, A, F) * (B, A, A, 1)
                # print("edge_cutoff:\n", edge_cutoff)
                edge_vec = F.mul(edge_vec, edge_cutoff)
        elif edge_mask is not None:
            # (B, A, A, F) * (B, A, A, 1)
            edge_mask = F.expand_dims(edge_mask, axis=-1)
            edge_vec = F.mul(edge_vec, F.expand_dims(edge_mask, axis=-1))
        else:
            edge_vec = edge_vec
        
        # (B, A, A, F) -> (B, F)
        agg = F.reduce_sum(edge_vec, self.axis)
        # (B, F) / (B, 1) -> (B, F)
        agg = F.div(agg, F.cast(num_edges, ms.float32))
        
        return agg
    
    def __str__(self):
        return "EdgeTensorMean<>"

_AGGREGATOR_BY_NAME = {
    agg.__name__: agg for agg in _AGGREGATOR_BY_KEY.values()}

def get_edge_aggregator(cls_name: Union[EdgeAggregator, str, dict],
                        dim: int,
                        axis: int = -2,
                        **kwargs,
                        ) -> EdgeAggregator:
    """get aggregator by name"""
    if cls_name is None or isinstance(cls_name, EdgeAggregator):
        return cls_name
    if isinstance(cls_name, dict):
        return get_edge_aggregator(**cls_name)
    if isinstance(cls_name, str):
        if cls_name.lower() == 'none':
            return None
        if cls_name.lower() in _AGGREGATOR_BY_KEY.keys():
            return _AGGREGATOR_BY_KEY[cls_name.lower()](dim=dim, axis=axis, **kwargs)
        if cls_name in _AGGREGATOR_BY_NAME.keys():
            return _AGGREGATOR_BY_NAME[cls_name](dim=dim, axis=axis, **kwargs)
        raise ValueError(
            "The Aggregator corresponding to '{}' was not found.".format(cls_name))
    raise TypeError(
        "Unsupported Aggregator type '{}'.".format(type(cls_name)))